#!/usr/bin/env python3
"""
数据库迁移脚本：为用户表添加balance字段
"""

import os
import sys
from sqlalchemy import create_engine, text
from config import Config

def migrate_add_balance():
    """添加balance字段到用户表"""
    try:
        # 创建数据库连接
        engine = create_engine(Config.SQLALCHEMY_DATABASE_URI)
        
        with engine.connect() as conn:
            # 检查balance字段是否已存在
            result = conn.execute(text("PRAGMA table_info(users)"))
            columns = [row[1] for row in result.fetchall()]
            
            if 'balance' in columns:
                print("✅ balance字段已存在，无需迁移")
                return True
            
            # 添加balance字段
            print("🔄 正在添加balance字段...")
            conn.execute(text("ALTER TABLE users ADD COLUMN balance NUMERIC(10, 2) DEFAULT 0.00 NOT NULL"))
            conn.commit()
            
            # 为现有用户设置默认余额
            print("🔄 为现有用户设置默认余额...")
            conn.execute(text("UPDATE users SET balance = 0.00 WHERE balance IS NULL"))
            conn.commit()
            
            print("✅ balance字段添加成功！")
            return True
            
    except Exception as e:
        print(f"❌ 迁移失败: {str(e)}")
        return False

def main():
    """主函数"""
    print("🚀 开始数据库迁移...")
    print("=" * 50)
    
    if migrate_add_balance():
        print("=" * 50)
        print("✅ 数据库迁移完成！")
        print("📝 用户表现在包含以下新功能：")
        print("   - balance字段：用户币余额")
        print("   - add_balance()方法：增加余额")
        print("   - subtract_balance()方法：减少余额")
        print("   - get_balance()方法：获取余额")
    else:
        print("=" * 50)
        print("❌ 数据库迁移失败！")
        sys.exit(1)

if __name__ == '__main__':
    main()